///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/optimizer.cpp
///
/// \brief   implementation file for the energy optimization function
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<vector>
#include<numeric>
#include<cassert>
#include<algorithm>
#include<cmath>
#include<utility>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/fqmc/optimizer.h>
#include<formic/fqmc/linear_method.h>
#include<formic/fqmc/bfgs.h>
#include<formic/fqmc/diis.h>
#include<formic/wfn/configuration.h>
#include<formic/fqmc/vmc_walker.h>
#include<formic/fqmc/energy.h>
#include<formic/fqmc/gradient.h>
#include<formic/fqmc/onepdm.h>
#include<formic/fqmc/s_squared.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/input/base.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/random/random.h>
#include<formic/output/output.h>
#include<formic/exception.h>
#include<formic/lapack/interface.h>
#include<formic/timing/timing.h>
#include<formic/fqmc/linear_method_classes.h>

namespace formic {

  namespace fqmc {

    struct OptUpdateCompare {

      bool operator()(const std::pair<char,double> & x, const std::pair<char,double> & y) {
        return x.second > y.second;
      }

    };

  }

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Optimize the variables of the supplied coefficient value object via the linear method.
///
/// \param[in]      userinp    the user's input options
/// \param[in,out]  sampler    object controlling sampling
/// \param[in]      co         a smart pointer to the coefficient object to optimize
/// \param[in]      update_gf  if true, the guiding function is updated each iteration
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::fqmc::optimize(const formic::InputBase & userinp,
                                              boost::shared_ptr<formic::fqmc::Sampler<S> > & sampler,
                                              const boost::shared_ptr<formic::CoeffObj<S> > & co,
                                              const bool update_gf) {

  struct local_funcs {

    static int n_real_vars(const boost::shared_ptr<formic::CoeffObj<double> > & co) {
      return co->n_der();
    }

    static int n_real_vars(const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > & co) {
      return 2 * co->n_der();
    }

    static void cplx_to_real(const size_t n, const double * const c, double * const r) {
      std::copy(c, c+n, r);
    }

    static void cplx_to_real(const size_t n, const std::complex<double> * const c, double * const r) {
      for (size_t i = 0; i < n; i++)
        r[i] = c[i].real();
      for (size_t i = 0; i < n; i++)
        r[i+n] = c[i].imag();
    }

  };

  // print a greeting
  if (formic::mpi::rank() == 0)
    formic::of << "Starting energy optimization." << std::endl << std::endl;

  // print the initial wavefunction variables
  if ( formic::mpi::rank() == 0 && userinp.get<bool>("opt_print_wfn") )
    co->print_variables();

  // create linear method updater
  formic::fqmc::LinMethodUpdater<S> lmu(userinp);

  // create bfgs updater
  formic::fqmc::BFGSUpdater<S> bfu(userinp);

  // create diis updater
  formic::fqmc::DIISUpdater<S> dsu(userinp);

  // create some work space
  std::vector<S> work_vec( ( formic::mpi::rank() == 0 ? co->n_der() : 1 ), formic::zero(S()));

  // create a vector to hold the history of real wavefunction parameters
  std::vector<boost::shared_ptr<std::vector<double> > > pmtr_hist;

  // create a vector to hold the history of energy gradients
  std::vector<boost::shared_ptr<std::vector<double> > > grad_hist;

  // sort out the order in which to apply the optimizations
  const char   lm_code = 'l';
  const char bfgs_code = 'b';
  const char diis_code = 'd';
  std::map<char,std::string> opt_code_str_map;
  opt_code_str_map[lm_code]   = "Linear Method";
  opt_code_str_map[bfgs_code] = "BFGS";
  opt_code_str_map[diis_code] = "DIIS";
  std::vector<std::pair<char,double> > opt_order;
  opt_order.push_back( std::pair<char,double>(  lm_code, userinp.get<double>("opt_lm_thresh"))   );
  opt_order.push_back( std::pair<char,double>(bfgs_code, userinp.get<double>("opt_bfgs_thresh")) );
  opt_order.push_back( std::pair<char,double>(diis_code, userinp.get<double>("opt_diis_thresh")) );
  std::sort(opt_order.begin(), opt_order.end(), formic::fqmc::OptUpdateCompare());
  if (formic::mpi::rank() == 0) {
    formic::of << boost::format("optimization order:") << std::endl;
    for (int i = 0; i < opt_order.size(); i++)
      formic::of << boost::format("%15.2e      %s") % opt_order.at(i).second % opt_code_str_map[opt_order.at(i).first] << std::endl;
    formic::of << std::endl << std::endl;
  }

  //// get the guiding function
  //boost::shared_ptr<formic::CoeffObj<S> > guide = ( gf ? gf : co->clone() );

  // clean up unused vector storage before starting the optimization iterations
  formic::vector_garbage_collect();

  // iteratively optimize the wavefunction
  bool converged = false;
  double last_energy = 0.0;
  std::string satisfied_criteria;
  for (int iter = 0; iter < userinp.get<int>("max_macro_loop"); iter++) {

    //// create an object to hold the starting configuration
    //formic::Configuration start_config(userinp);

    //// create an object to hold the starting state of the random number generator
    //formic::LaggedFibonacci<> start_rng;

    // if requested, update guiding function and burn in again
    if (update_gf)
      sampler->set_gf(co->clone());

//    // burn in the markov chain
//    if ( !userinp.get<bool>("exact_sampling") ) {
//
//      // prepare a walker
//      formic::VMC_Walker<S> walker(userinp);
//      walker.gf() = guide;
//
//      // burn in
//      formic::sample(userinp, userinp.get<int>("opt_burn_len"), walker);
//
//      // save the configuration and the state of the random number generator
//      start_config = walker.config();
//      start_rng = walker.lfg();
//
//    }

    // print the initial wavefunction's natural orbitals
    if ( iter == 0 && userinp.get<bool>("opt_print_wfn") )
      formic::fqmc::print_nos(userinp, sampler, co);

    // print the initial wave function's S^2 expectation value
    if ( iter == 0 ) {
      const S s2 = formic::fqmc::s_squared(userinp, sampler, co);
      if (formic::mpi::rank() == 0)
        formic::of << boost::format("S^2 = %12.6f    %s    %12.6f i")
                        % formic::real(s2)
                        % ( formic::imag(s2) >= 0.0 ? "+" : "-" )
                        % std::abs(formic::imag(s2))
                   << std::endl
                   << std::endl;
    }

    // compute energy and standard deviation of the initial wavefunction for this iteration
    if (formic::mpi::rank() == 0)
      formic::of << boost::format("computing energy of initial wavefunction...") << std::endl << std::endl;
    double init_sdev = 0.0;
    double init_serr = 0.0;
    const double init_energy = formic::fqmc::energy(userinp, sampler, co, init_sdev, init_serr);

    // get space to hold this iteration's gradient and throw out old gradients we no longer need
    const int hist_len = std::max(userinp.get<int>("bfgs_hist_len"), userinp.get<int>("diis_hist_len"));
    boost::shared_ptr<std::vector<double> > grad_ptr;
    if ( formic::mpi::rank() == 0 ) {
      if ( grad_hist.size() > 0 && grad_hist.size() > hist_len ) {
        grad_ptr = *grad_hist.rbegin();
        grad_hist.pop_back();
      } else {
        boost::shared_ptr<std::vector<double> > temp_ptr( new std::vector<double>(std::max(local_funcs::n_real_vars(co), 1), 0.0) );
        grad_ptr = temp_ptr;
      }
      grad_hist.insert(grad_hist.begin(), grad_ptr);
    } else {
      if ( grad_hist.size() == 0 ) {
        boost::shared_ptr<std::vector<double> > temp_ptr( new std::vector<double>(std::max(local_funcs::n_real_vars(co), 1), 0.0) );
        grad_hist.insert(grad_hist.begin(), temp_ptr);
      }
      grad_ptr = grad_hist.at(0);
    }

    // compute the gradient of the initial wavefunction for this iteration
    double * const grad = &grad_ptr->at(0);
    formic::fqmc::gradient(userinp, sampler, co, grad, &init_energy);

    // compute the gradient norm
    double grad_norm = 0.0;
    if ( userinp.get<bool>("overlap_grad_norm") ) {
      formic::fqmc::OvlpOpLM<S> ovl_op(userinp, sampler, co);
      formic::Vector<S,int> g_vec(1+co->n_der(), formic::zero(S()));
      formic::Vector<S,int> sg_vec(1+co->n_der(), formic::zero(S()));
      if ( formic::mpi::rank() == 0 )
        formic::compress_to_complex(co->n_der(), grad, &g_vec.at(1));
      ovl_op.operate(g_vec.size(), &g_vec.at(0), &sg_vec.at(0), formic::mpi::rank(), 0);
      if ( formic::mpi::rank() == 0 )
        grad_norm = std::abs( std::sqrt( formic::xdotc(g_vec.size(), &g_vec.at(0), 1, &sg_vec.at(0), 1) ) );
    } else {
      if ( formic::mpi::rank() == 0 )
        grad_norm = std::abs( std::sqrt( formic::xdotc(grad_ptr->size(), grad, 1, grad, 1) ) );
    }
    formic::mpi::bcast(grad_norm);

    // print the status of the optimization
    if (formic::mpi::rank() == 0)
      formic::of << boost::format("starting energy optimization step %4i:     grad norm = %.4e     energy = %20.12f") % (iter+1) % grad_norm % init_energy << std::endl << std::endl;

    // if the gradient is small enough, stop optimizing
    if ( grad_norm < userinp.get<double>("opt_gr_thresh") ) {
      converged = true;
      satisfied_criteria = "gradient";
    // if the energy difference is small enough, stop optimizing
    } else if ( iter > 0 && std::abs( init_energy - last_energy ) < userinp.get<double>("opt_en_thresh") ) {
      converged = true;
      satisfied_criteria = "energy";
    }
    formic::mpi::bcast(converged);

    //// check gradient using numeric derivatives
    //if ( iter > 20 && ( iter + 1 == userinp.get<int>("max_macro_loop") || converged ) ) {
    //  formic::of << boost::format("checking gradient:") << std::endl;
    //  std::vector<S> temp(co->n_der(), formic::zero(S()));
    //  co->to_vector(&temp.at(0));
    //  for (int i = 0; i < co->n_der(); i++) {
    //    const S old_temp = temp.at(i);
    //    //const S dx = 1.0e-8 * formic::unity(S()) * std::abs(temp.at(i));
    //    const S dx = 1.0e-8;
    //    temp.at(i) += dx;
    //    co->from_vector(&temp.at(0));
    //    double sdev = 0.0;
    //    const S grd_elem = ( formic::fqmc::energy(userinp, 0, userinp.get<int>("opt_samp_len"), start_config, start_rng, co, co, sdev) - init_energy ) / dx;
    //    formic::of << boost::format("gradient element %4i:      %20.12f  numeric      %20.12f  analytic     %20.12f  error")
    //                  % i % grd_elem % grad[i] % std::abs(grd_elem-grad[i])
    //               << std::endl;
    //    temp.at(i) = old_temp;
    //  }
    //  formic::of << std::endl;
    //  co->from_vector(&temp.at(0));
    //  //throw formic::Exception("done checking gradient");
    //}

    if (converged)
      break;

    // save the energy
    last_energy = init_energy;

    // save the current wavefunction variables
    if ( formic::mpi::rank() == 0 ) {
      if ( pmtr_hist.size() > 0 && pmtr_hist.size() > hist_len ) {
        boost::shared_ptr<std::vector<double> > temp_ptr = *pmtr_hist.rbegin();
        pmtr_hist.pop_back();
        pmtr_hist.insert(pmtr_hist.begin(), temp_ptr);
      } else {
        boost::shared_ptr<std::vector<double> > temp_ptr( new std::vector<double>(std::max(local_funcs::n_real_vars(co), 1), 0.0) );
        pmtr_hist.insert(pmtr_hist.begin(), temp_ptr);
      }
      co->to_vector(&work_vec.at(0));
      local_funcs::cplx_to_real(work_vec.size(), &work_vec.at(0), &(**pmtr_hist.begin()).at(0));
    }

    // decide which type of update to do
    char opt_code = opt_order.at(0).first;
    for (int i = 1; i < opt_order.size() && grad_norm < opt_order.at(i).second; i++)
      opt_code = opt_order.at(i).first;
    formic::mpi::bcast(opt_code);
    if (formic::mpi::rank() == 0)
      formic::of << boost::format("performing a %s update") % opt_code_str_map[opt_code] << std::endl << std::endl;

    // perform the update
    bool rejected = false;
    if ( opt_code == bfgs_code ) {
      rejected = bfu.update(userinp, pmtr_hist, grad_hist, sampler, co);
    } else if ( opt_code == lm_code ) {
      rejected = lmu.update(userinp, sampler, co);
    } else if ( opt_code == diis_code ) {
      rejected = dsu.update(userinp, pmtr_hist, grad_hist, sampler, co);
    } else {
      throw formic::Exception("unknown opt_code \"%s\"") % opt_code;
    }
    formic::mpi::bcast(rejected);

    // if the update was rejected, remove the current wavefunction and gradient from the history
    if ( formic::mpi::rank() == 0 && rejected ) {
      pmtr_hist.erase(pmtr_hist.begin());
      grad_hist.erase(grad_hist.begin());
    }

    // print the new wavefunction variables
    if ( formic::mpi::rank() == 0 && userinp.get<bool>("opt_print_wfn") )
      co->print_variables();

    // print the wavefunction's natural orbitals
    if ( userinp.get<bool>("opt_print_wfn") )
      formic::fqmc::print_nos(userinp, sampler, co);

    // print the wave function's S^2 expectation value
    {
      const S s2 = formic::fqmc::s_squared(userinp, sampler, co);
      if (formic::mpi::rank() == 0)
        formic::of << boost::format("S^2 = %12.6f    %s    %12.6f i")
                        % formic::real(s2)
                        % ( formic::imag(s2) >= 0.0 ? "+" : "-" )
                        % std::abs(formic::imag(s2))
                   << std::endl
                   << std::endl;
    }

//    // if we were not supplied with a guiding function, update the guiding function
//    if ( !gf ) {
//      guide.reset();
//      guide = co->clone();
//    }

//    // reset vector storage
//    formic::vector_garbage_collect();

  }

  // clean up unused vector storage now that optimization iterations are finished
  formic::vector_garbage_collect();

  // print final outcome
  if ( formic::mpi::rank() == 0 && converged )
    formic::of << boost::format("optimization converged based on %s criteria") % satisfied_criteria << std::endl << std::endl;
  else if ( formic::mpi::rank() == 0 )
    formic::of << boost::format("optimization did not converge") << std::endl << std::endl;

//  // compute final energy
//  if (formic::mpi::rank() == 0)
//    formic::of << boost::format("computing energy of final wavefunction...") << std::endl << std::endl;
//  formic::fqmc::energy(userinp, co);

}

template void formic::fqmc::optimize(const formic::InputBase & userinp,
                                     boost::shared_ptr<formic::fqmc::Sampler<double> > & sampler,
                                     const boost::shared_ptr<formic::CoeffObj<double> > & co,
                                     const bool update_gf);
template void formic::fqmc::optimize(const formic::InputBase & userinp,
                                     boost::shared_ptr<formic::fqmc::Sampler<std::complex<double> > > & sampler,
                                     const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > & co,
                                     const bool update_gf);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Optimize the variables of the supplied coefficient value object via the linear method.
///
/// \param[in]      userinp  the user's input options
/// \param[in]      co       a smart pointer to the coefficient object to optimize
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::fqmc::optimize(const formic::InputBase & userinp, const boost::shared_ptr<formic::CoeffObj<S> > & co) {

  // create an empty pointer to feed in as the guiding function, which will tell the optimizer to use co as its own guiding function
  boost::shared_ptr<formic::CoeffObj<S> > gf;

  // create the object that will control sampling
  boost::shared_ptr<formic::fqmc::Sampler<S> > sampler( new formic::fqmc::Sampler<S>(userinp,
                                                                                     userinp.get<int>("opt_burn_len"),
                                                                                     userinp.get<int>("opt_samp_len"),
                                                                                     formic::Configuration(userinp),
                                                                                     formic::LaggedFibonacci<>(),
                                                                                     co->clone()) );

  // perform the optimization, updating the guiding function each iteration
  const bool update_gf = true;
  formic::fqmc::optimize(userinp, sampler, co, update_gf);

}

template void formic::fqmc::optimize(const formic::InputBase &, const boost::shared_ptr<formic::CoeffObj<double> > &);
template void formic::fqmc::optimize(const formic::InputBase &, const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > &);

